from cryptography import x509
from cryptography.x509.oid import NameOID, ExtendedKeyUsageOID
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.backends import default_backend
import datetime
import os

# 创建目录
os.makedirs("certificates/root-ca", exist_ok=True)
os.makedirs("certificates/intermediate-ca", exist_ok=True)
os.makedirs("certificates/server", exist_ok=True)


# 生成RSA私钥
def generate_private_key(path, password=None):
    private_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
        backend=default_backend()
    )

    # 保存私钥
    encryption = serialization.NoEncryption()
    if password:
        encryption = serialization.BestAvailableEncryption(password.encode())

    with open(path, "wb") as f:
        f.write(private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=encryption
        ))

    return private_key


# 生成自签名根证书
def generate_root_certificate(private_key, subject, path):
    # 创建证书构建器
    builder = x509.CertificateBuilder()
    builder = builder.subject_name(subject)
    builder = builder.issuer_name(subject)  # 自签名，颁发者和主题相同
    builder = builder.not_valid_before(datetime.datetime.utcnow())
    builder = builder.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=3650))
    builder = builder.serial_number(x509.random_serial_number())
    builder = builder.public_key(private_key.public_key())

    # 添加扩展
    builder = builder.add_extension(
        x509.BasicConstraints(ca=True, path_length=None),
        critical=True,
    )
    builder = builder.add_extension(
        x509.KeyUsage(
            digital_signature=True,
            content_commitment=False,
            key_encipherment=False,
            data_encipherment=False,
            key_agreement=False,
            key_cert_sign=True,
            crl_sign=True,
            encipher_only=False,
            decipher_only=False
        ),
        critical=True,
    )
    builder = builder.add_extension(
        x509.SubjectKeyIdentifier.from_public_key(private_key.public_key()),
        critical=False,
    )

    # 签名证书
    certificate = builder.sign(
        private_key=private_key,
        algorithm=hashes.SHA256(),
        backend=default_backend()
    )

    # 保存证书
    with open(path, "wb") as f:
        f.write(certificate.public_bytes(serialization.Encoding.PEM))

    return certificate


# 生成中间CA证书
def generate_intermediate_certificate(
        intermediate_key, intermediate_subject, issuer_cert, issuer_key, path
):
    builder = x509.CertificateBuilder()
    builder = builder.subject_name(intermediate_subject)
    builder = builder.issuer_name(issuer_cert.subject)
    builder = builder.not_valid_before(datetime.datetime.utcnow())
    builder = builder.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=1825))
    builder = builder.serial_number(x509.random_serial_number())
    builder = builder.public_key(intermediate_key.public_key())

    # 添加扩展
    builder = builder.add_extension(
        x509.BasicConstraints(ca=True, path_length=0),
        critical=True,
    )
    builder = builder.add_extension(
        x509.KeyUsage(
            digital_signature=True,
            content_commitment=False,
            key_encipherment=False,
            data_encipherment=False,
            key_agreement=False,
            key_cert_sign=True,
            crl_sign=True,
            encipher_only=False,
            decipher_only=False
        ),
        critical=True,
    )
    builder = builder.add_extension(
        x509.SubjectKeyIdentifier.from_public_key(intermediate_key.public_key()),
        critical=False,
    )
    builder = builder.add_extension(
        x509.AuthorityKeyIdentifier.from_issuer_public_key(issuer_key.public_key()),
        critical=False,
    )

    # 签名证书
    certificate = builder.sign(
        private_key=issuer_key,
        algorithm=hashes.SHA256(),
        backend=default_backend()
    )

    # 保存证书
    with open(path, "wb") as f:
        f.write(certificate.public_bytes(serialization.Encoding.PEM))

    return certificate


# 生成服务器证书
def generate_server_certificate(
        server_key, server_subject, dns_names, issuer_cert, issuer_key, path
):
    builder = x509.CertificateBuilder()
    builder = builder.subject_name(server_subject)
    builder = builder.issuer_name(issuer_cert.subject)
    builder = builder.not_valid_before(datetime.datetime.utcnow())
    builder = builder.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=365))
    builder = builder.serial_number(x509.random_serial_number())
    builder = builder.public_key(server_key.public_key())

    # 添加扩展
    builder = builder.add_extension(
        x509.BasicConstraints(ca=False, path_length=None),
        critical=True,
    )
    builder = builder.add_extension(
        x509.KeyUsage(
            digital_signature=True,
            content_commitment=False,
            key_encipherment=True,
            data_encipherment=False,
            key_agreement=False,
            key_cert_sign=False,
            crl_sign=False,
            encipher_only=False,
            decipher_only=False
        ),
        critical=True,
    )
    builder = builder.add_extension(
        x509.ExtendedKeyUsage([ExtendedKeyUsageOID.SERVER_AUTH]),
        critical=False,
    )
    builder = builder.add_extension(
        x509.SubjectAlternativeName([x509.DNSName(dns) for dns in dns_names]),
        critical=False,
    )
    builder = builder.add_extension(
        x509.SubjectKeyIdentifier.from_public_key(server_key.public_key()),
        critical=False,
    )
    builder = builder.add_extension(
        x509.AuthorityKeyIdentifier.from_issuer_public_key(issuer_key.public_key()),
        critical=False,
    )

    # 签名证书
    certificate = builder.sign(
        private_key=issuer_key,
        algorithm=hashes.SHA256(),
        backend=default_backend()
    )

    # 保存证书
    with open(path, "wb") as f:
        f.write(certificate.public_bytes(serialization.Encoding.PEM))

    return certificate


# 主函数：生成完整证书链
def main():
    # 定义证书主体信息
    root_subject = x509.Name([
        x509.NameAttribute(NameOID.COUNTRY_NAME, "CN"),
        x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Shanghai"),
        x509.NameAttribute(NameOID.LOCALITY_NAME, "Shanghai"),
        x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Demo Root CA"),
        x509.NameAttribute(NameOID.COMMON_NAME, "demo-root-ca"),
    ])

    intermediate_subject = x509.Name([
        x509.NameAttribute(NameOID.COUNTRY_NAME, "CN"),
        x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Shanghai"),
        x509.NameAttribute(NameOID.LOCALITY_NAME, "Shanghai"),
        x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Demo Intermediate CA"),
        x509.NameAttribute(NameOID.COMMON_NAME, "demo-intermediate-ca"),
    ])

    server_subject = x509.Name([
        x509.NameAttribute(NameOID.COUNTRY_NAME, "CN"),
        x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Shanghai"),
        x509.NameAttribute(NameOID.LOCALITY_NAME, "Shanghai"),
        x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Demo Server"),
        x509.NameAttribute(NameOID.COMMON_NAME, "localhost"),
    ])

    # 生成根CA私钥和证书
    print("生成根CA证书...")
    root_key = generate_private_key("certificates/root-ca/root.key.pem")
    root_cert = generate_root_certificate(
        root_key, root_subject, "certificates/root-ca/root.cert.pem"
    )

    # 生成中间CA私钥和证书
    print("生成中间CA证书...")
    intermediate_key = generate_private_key("certificates/intermediate-ca/intermediate.key.pem")
    intermediate_cert = generate_intermediate_certificate(
        intermediate_key, intermediate_subject, root_cert, root_key,
        "certificates/intermediate-ca/intermediate.cert.pem"
    )

    # 生成服务器私钥和证书
    print("生成服务器证书...")
    server_key = generate_private_key("certificates/server/server.key.pem")
    server_cert = generate_server_certificate(
        server_key, server_subject, ["localhost", "127.0.0.1"],
        intermediate_cert, intermediate_key, "certificates/server/server.cert.pem"
    )

    # 创建证书链文件
    with open("certificates/server/chain.cer", "wb") as f:
        f.write(server_cert.public_bytes(serialization.Encoding.PEM))
        f.write(intermediate_cert.public_bytes(serialization.Encoding.PEM))
        f.write(root_cert.public_bytes(serialization.Encoding.PEM))

    print("证书链生成完成！")
    print("根CA证书: certificates/root-ca/root.cert.pem")
    print("中间CA证书: certificates/intermediate-ca/intermediate.cert.pem")
    print("服务器证书: certificates/server/server.cert.pem")
    print("完整证书链: certificates/server/chain.cer")


if __name__ == "__main__":
    main()